import os
import glob
import argparse
from pathlib import Path
import gmalg

def verify(msk_s_file, mpk_s_file, sk_s_file, uid_file, msg_file, h_file, s_file):
    hid_s = b"\x01"
    mpk_s_temp = read_binary_file_(mpk_s_file)
    mpk_s = bytes.fromhex("04") + mpk_s_temp
    msk_s = read_binary_file_(msk_s_file)
    sk_s_temp = read_binary_file_(sk_s_file)
    sk_s = bytes.fromhex("04") + sk_s_temp
    uid = read_binary_file_(uid_file)
    msg = read_binary_file_(msg_file)
    h = read_binary_file_(h_file)
    s = read_binary_file_(s_file)

    kgc = gmalg.SM9KGC(hid_s=hid_s, msk_s=msk_s, mpk_s=mpk_s)
    
    sk_s_creat = kgc.generate_sk_sign(uid)
    
    #print(sk_s.hex())
    #print(sk_s_creat.hex())
    assert sk_s == sk_s_creat, "验证失败" 
    
    sm9 = gmalg.SM9(hid_s=hid_s, mpk_s=mpk_s, sk_s=sk_s, uid=uid)
    
    #message = b"Chinese IBS standard"
    #h, S = sm9.sign(msg)
    
    #print(h.hex())
    #print(S.hex())
    
    assert sm9.verify(msg, h, s) == True, "验证失败"


def read_binary_file_(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data


def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['sign', 'verify'], default='sign', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode
    print (f"模式:{mode}\n")

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    if mode == 'sign' :
        for i in range(20):
            # 读取文件内容
            r_file = os.path.join(directory, f'r_{i}')
            id_file = os.path.join(directory, f'id')
            dsA_file = os.path.join(directory, f'dsA')
            ks_file = os.path.join(directory, f'ks')
            Ppub_s_file = os.path.join(directory, f'Ppub_s')
            msg_file = os.path.join(directory, f'msg')
            h_file = os.path.join(directory, f'h_{i}')
            s_file = os.path.join(directory, f's_{i}')

            verify(ks_file, Ppub_s_file, dsA_file, id_file, msg_file, h_file, s_file)
            r_content = read_binary_file(r_file)
            id_content = read_binary_file(id_file)
            Ppub_s_content = read_binary_file(Ppub_s_file)
            dsA_content = read_binary_file(dsA_file)
            msg_content = read_binary_file(msg_file)
            h_content = read_binary_file(h_file)
            s_content = read_binary_file(s_file)

            # 添加到合并内容
            merged_content += f"随机数=\n{r_content}\n"
            merged_content += f"用户私钥=\n{dsA_content}\n"
            merged_content += f"公钥=\n{Ppub_s_content}\n"
            merged_content += f"用户标识=\n{id_content}\n"
            merged_content += f"签名数据=\n{msg_content}\n"
            merged_content += f"签名h=\n{h_content}\n"
            merged_content += f"签名s=\n{s_content}\n"
            merged_content += "\n"
    else: 
        dsA_file = os.path.join(directory, f'dsA')
        ks_file = os.path.join(directory, f'ks')
        Ppub_s_file = os.path.join(directory, f'Ppub_s')
        id_file = os.path.join(directory, f'id')
        h_file = os.path.join(directory, f'h')
        s_file = os.path.join(directory, f's')
        msg_file = os.path.join(directory, f'msg')

        verify(ks_file, Ppub_s_file, dsA_file, id_file, msg_file, h_file, s_file)
        Ppub_s_content = read_binary_file(Ppub_s_file)
        id_content = read_binary_file(id_file)
        h_content = read_binary_file(h_file)
        s_content = read_binary_file(s_file)
        msg_content = read_binary_file(msg_file)

        # 添加到合并内容
        merged_content += f"公钥=\n{Ppub_s_content}\n"
        merged_content += f"用户标识=\n{id_content}\n"
        merged_content += f"签名h=\n{h_content}\n"
        merged_content += f"签名s=\n{s_content}\n"
        merged_content += f"签名数据=\n{msg_content}\n"
        merged_content += "\n"

    # 写入到txt文件
    if mode == 'sign':
        output_file = 'sm9签名性能.txt'
    else:
        output_file = 'sm9验签性能.txt'

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
